#include <string>
#include <cstring>
#include <iostream>
#include <map>
#include <vector>
#include <sstream>
#include <iomanip>
#include <stdexcept>
#include <memory>
#include <algorithm>
#include <cctype>
#include <utility>

#include <dlfcn.h>
#include <stdlib.h>
#include <unistd.h>
#include <cxxabi.h>
#include <link.h>

#include <cuda_runtime.h>

#define HOOK_LOG_PREFIX "[HOOK] "

#define CHECK_CUDA_ERROR(call) \
    do { \
        cudaError_t err = (call); \
        if (err != cudaSuccess) { \
            std::cerr << HOOK_LOG_PREFIX << "CUDA error at " << __FILE__ << ":" << __LINE__ << ": " \
                      << cudaGetErrorString(err) << std::endl; \
        } \
    } while(0)

enum class DataType {UNREC = -1, INT, FLOAT, BOOL, PTR, DOUBLE, SIZE_T, UINT};

using TypeInfo = std::pair<size_t, DataType>;

std::string format_void_ptr(int index, void* ptr, DataType type) {
    std::ostringstream oss;
    oss << "(" << index << ",";
    switch (type) {
        case DataType::INT:
            oss << "int," << *static_cast<int*>(ptr);
            break;
        case DataType::FLOAT:
            oss << "float," << std::fixed << std::setprecision(2) << *static_cast<float*>(ptr);
            break;
        case DataType::BOOL:
            oss << "bool," << std::boolalpha << *static_cast<bool*>(ptr);
            break;
        case DataType::PTR:
            oss << "ptr," << *static_cast<size_t *>(ptr);
            break;
        case DataType::DOUBLE:
            oss << "double," << std::fixed << std::setprecision(2) << *static_cast<double*>(ptr);
            break;
        case DataType::SIZE_T:
            oss << "size_t," << *static_cast<size_t*>(ptr);
            break;
        case DataType::UINT:
            oss << "uint," << *static_cast<unsigned int*>(ptr);
            break;
        default:
            oss << "??";
            break;
    }
    oss << ")";
    return oss.str();
}

std::string format_void_array(void** data_array, const std::vector<TypeInfo>& type_info) {
    if(type_info.empty()) return "??";
    std::string result;
    for (size_t i = 0; i < type_info.size(); ++i) {
        const auto& [index, type] = type_info[i];
        if (!result.empty()) {
            result += ",";
        }
        if(i != 0 && i % 4 == 0){
            result += "\n  ";
        }
        result += format_void_ptr(index, data_array[index], type);
    }
    return result;
}

std::string trim(const std::string& str) {
    auto start = str.find_first_not_of(" ");
    auto end = str.find_last_not_of(" ");
    return (start == std::string::npos) ? "" : str.substr(start, end - start + 1);
}

std::vector<TypeInfo> parseFuncParams(const std::string& args_decl) {
    auto paramStrToTypeInfo = [](size_t &index, const std::string &param_str) -> TypeInfo {
        if (param_str.find('<') != std::string::npos || param_str.find('>') != std::string::npos) {
            return {index++, DataType::UNREC};
        }
        if (param_str.find('*') != std::string::npos) {
            return {index++, DataType::PTR};
        }
        if (param_str == "int" || param_str == "long" || param_str == "signed") {
            return {index++, DataType::INT};
        } else if (param_str == "float") {
            return {index++, DataType::FLOAT};
        } else if (param_str == "double") {
            return {index++, DataType::DOUBLE};
        } else if (param_str == "bool") {
            return {index++, DataType::BOOL};
        }else if (param_str == "unsigned int" || param_str == "unsigned" || param_str == "uint32_t") {
            return {index++, DataType::UINT};
        } else if (param_str == "size_t" || param_str == "uint64_t") {
            return {index++, DataType::SIZE_T};
        } else {
            return {index++, DataType::UNREC};
        }
    };
    std::vector<TypeInfo> result;
    size_t start_paren = args_decl.find('(');
    size_t end_paren = args_decl.rfind(')');
    
    if (start_paren == std::string::npos || end_paren == std::string::npos || 
        start_paren >= end_paren) {
        return result; // invalid func dec
    }
    std::string params_str = args_decl.substr(start_paren + 1, end_paren - start_paren - 1);
    size_t index = 0;
    
    // 分割参数
    int depth = 0;  // 模板嵌套深度
    size_t param_start = 0;
    size_t start = 0, end = params_str.size();

    for (size_t i = start; i < end; ++i) {
        char c = params_str[i];
        if (c == '<') depth++;
        else if (c == '>') depth--;
        else if (c == ',' && depth == 0) {
            // 提取当前参数并去除首尾空白
            size_t pstart = param_start;
            size_t pend = i;
            
            result.emplace_back(paramStrToTypeInfo(index, trim(params_str.substr(pstart, pend - pstart))));
            param_start = i + 1;  // 下一个参数起始位置
        }
    }
    
    // 添加最后一个参数
    if (param_start < end) {
        size_t pstart = param_start;
        size_t pend = end;
        result.emplace_back(paramStrToTypeInfo(index, trim(params_str.substr(pstart, pend - pstart))));
    }
    return result;
}

void extract_name_and_args(const std::string& decl, std::string &func_name, std::string &args_decl) {
    // 步骤1: 正确找到函数参数列表起始位置
    int depth = 0;
    size_t param_start = std::string::npos;
    
    for (auto it = decl.rbegin(); it != decl.rend(); ++it) {
        char c = *it;
        if (c == ')') {
            depth++;
        } else if (c == '(') {
            depth--;
            if (depth == 0) {
                param_start = decl.rend() - it - 1;
                break;
            }
        }
    }
    
    if (param_start == std::string::npos) {
        func_name = "??", args_decl = "??";
        return;
    }

    // 步骤2: 提取前缀部分
    std::string prefix = decl.substr(0, param_start);

    // 步骤3: 去除尾部空白
    auto end_it = std::find_if_not(prefix.rbegin(), prefix.rend(), 
        [](char c) { return std::isspace(c); });
    
    if (end_it != prefix.rend()) {
        prefix.erase(end_it.base(), prefix.end());
    } else {
        prefix.clear();
    }

    if (prefix.empty()) {
        func_name = "??", args_decl = "??";
        return;
    }

    // 步骤4: 反向扫描定位函数名起始
    int bracket_count = 0;   // <> 计数
    int paren_count = 0;     // () 计数
    
    auto name_start = std::find_if(prefix.rbegin(), prefix.rend(), 
        [&](char c) {
            // 更新括号计数
            if (c == '>') bracket_count++;
            else if (c == ')') paren_count++;
            else if (c == '<' && bracket_count > 0) bracket_count--;
            else if (c == '(' && paren_count > 0) paren_count--;
            
            // 在模板外且无括号嵌套时遇到空格停止
            return (bracket_count == 0 && paren_count == 0) && std::isspace(c);
        });

    // 步骤5: 提取函数名部分
    func_name = (name_start == prefix.rend()) ? prefix : std::string(name_start.base(), prefix.end());
    args_decl = decl.substr(param_start, std::string::npos);
}

std::string demangle(const std::string& mangle_name) {
    int status = 0;
    std::string demangled = abi::__cxa_demangle(mangle_name.c_str(), nullptr, nullptr, &status);
    if (status != 0) {
        std::cerr << HOOK_LOG_PREFIX << "demangle err " << mangle_name << std::endl;
        return "??";
    }
    return demangled;
}

std::string DataTypeToStr(DataType dataType){
    std::string dataTypeStr;
    switch (dataType) {
        case DataType::INT:
            dataTypeStr = "int";;
            break;
        case DataType::FLOAT:
            dataTypeStr = "float";
            break;
        case DataType::BOOL:
            dataTypeStr = "bool";
            break;
        case DataType::PTR:
            dataTypeStr = "ptr";
            break;
        case DataType::DOUBLE:
            dataTypeStr = "double";
            break;
        case DataType::SIZE_T:
            dataTypeStr = "size_t";
            break;
        case DataType::UINT:
            dataTypeStr = "uint";
            break;
        default:
            dataTypeStr = "??";
            break;
    }
    return dataTypeStr;
}

void get_params(const std::string &mangle_func_dec, void **params, std::string &func_name, std::string &params_str){
    static std::map<std::string, std::pair<std::string, std::vector<TypeInfo>>> func_infos = {
        {"??", {"??", {}}}
        //you can add something
    };
    if (func_infos.find(mangle_func_dec) == func_infos.end()) {
        std::string demangle_func_dec = demangle(mangle_func_dec), func_name, args_decl;
        extract_name_and_args(demangle_func_dec, func_name, args_decl);
        func_infos.emplace(mangle_func_dec, std::make_pair(func_name, parseFuncParams(args_decl)));
    }
    func_name = func_infos[mangle_func_dec].first;
    params_str = format_void_array(params, func_infos[mangle_func_dec].second);
}

std::string load_func_name(const void *func){
    std::string func_name = "??";
    Dl_info info;
    if (dladdr(func, &info)) {
        if(info.dli_sname) func_name = info.dli_sname;
        // std::cout << HOOK_LOG_PREFIX << "File path: " << info.dli_fname << std::endl;
    } else {
        std::cerr << HOOK_LOG_PREFIX << " Failed to resolve function name." << std::endl;
    }
    return func_name;
}

bool isInContext(const char *func_name){
    const char* env_var = std::getenv("CUDA_API_HOOK_CTX");
    bool cuda_api_hook_ctx_on = false;
    if (env_var) {
        if (std::strcmp(env_var, "ON") == 0) {
            cuda_api_hook_ctx_on = true;
        } else {
            cuda_api_hook_ctx_on = false;
        }
    }
    if(env_var && !cuda_api_hook_ctx_on){
        if(std::strcmp(env_var, func_name) == 0) {
            cuda_api_hook_ctx_on = true;
        }
    }
    return cuda_api_hook_ctx_on;
}

template<typename T>
T load_func(const std::string &func_name, const std::string &so_name){
    T func_ptr = nullptr;
    static char buf[128];
    // load so
    void *handle = dlopen(so_name.c_str(), RTLD_LAZY | RTLD_GLOBAL);
    if (!handle) {
        std::cerr << HOOK_LOG_PREFIX << " Error in dlopen: " << dlerror() << std::endl;
        return nullptr;
    }
    func_ptr = reinterpret_cast<T>(dlsym(handle, func_name.c_str()));
    if (!func_ptr) {
        std::cerr << HOOK_LOG_PREFIX << " Error in dlsym: " << dlerror() << std::endl;
        return nullptr;
    }
    return func_ptr;
}

static int shared_lib_callback(struct dl_phdr_info* info, size_t, void* data)
{   
    if (info->dlpi_name == nullptr || info->dlpi_name[0] == '\0') return 0;
    
    const std::string lib_path(info->dlpi_name);
    
    auto last_slash_pos = lib_path.find_last_of('/');
    std::string lib_name = (last_slash_pos == std::string::npos) 
                           ? lib_path 
                           : lib_path.substr(last_slash_pos + 1);
    
    auto& ctx = *static_cast<std::pair<const std::string*, std::vector<std::string>*>*>(data);
    auto& target_lib = *(ctx.first);
    auto& results = *(ctx.second);
    if (lib_name.find(target_lib) != std::string::npos) {
        char absolute_path[PATH_MAX];
        if (realpath(lib_path.c_str(), absolute_path)) {
            results.push_back(absolute_path);
        } else {
            results.push_back(lib_path);
        }
    }
    return 0;
}

std::string getEnv(const std::string& envName) {
    const char* value = std::getenv(envName.c_str());
    if (value != nullptr) {
        return std::string(value);
    }
    return "";
}

std::string get_shared_library_path(const std::string& library_name)
{
    std::vector<std::string> results;
    std::string shared_library_path = library_name;
    auto ctx = std::pair(&library_name, &results);
    dl_iterate_phdr(shared_lib_callback, &ctx);
    if (!results.empty()) {
        shared_library_path = results[0];
        std::cout << HOOK_LOG_PREFIX << "Found " << results.size() << " instance(s) of " << library_name << ", use " << shared_library_path << std::endl;
    } else {
        std::cerr << HOOK_LOG_PREFIX << shared_library_path << " not found in loaded libraries, will be used system path" << std::endl;
    }
    return shared_library_path;
}

typedef cudaError_t (*cudaMalloc_t)(void **devPtr, size_t size);
typedef cudaError_t (*cudaFree_t)(void *devPtr);
typedef cudaError_t (*cudaLaunchKernel_t)(const void *func, dim3 gridDim, dim3 blockDim, void **args, size_t sharedMem, cudaStream_t stream);

extern "C" {
    
    //cudaMalloc hook
    cudaError_t cudaMalloc(void **devPtr, size_t size){
        // real target func
        static cudaMalloc_t real_cudaMalloc = nullptr;

        if (!real_cudaMalloc) {
            auto lib_path = get_shared_library_path("libcudart.so");
            real_cudaMalloc = load_func<cudaMalloc_t>("cudaMalloc", lib_path);
            if(!real_cudaMalloc) return cudaErrorInvalidValue;
        }
        cudaError_t err = real_cudaMalloc(devPtr, size);
        if(isInContext("cudaMalloc")){
            std::cout << HOOK_LOG_PREFIX << "  call cudaMalloc (" <<  *(size_t *)devPtr << "," << size << ")"<< std::endl;
        }
        return err;
    }

    //cudaFree hook
    cudaError_t cudaFree(void *devPtr){
        // real target func
        static cudaFree_t real_cudaFree = nullptr;

        if (!real_cudaFree) {
            auto lib_path = get_shared_library_path("libcudart.so");
            real_cudaFree = load_func<cudaFree_t>("cudaFree", lib_path);
            if(!real_cudaFree) return cudaErrorInvalidValue;
        }
        cudaError_t err = real_cudaFree(devPtr);
        if(isInContext("cudaFree")){
            //do something
            std::cout << HOOK_LOG_PREFIX << "  call cudaFree " << size_t(devPtr) << std::endl;
        }
        return err;
    }

    //cudaLaunchKernel hook
    cudaError_t cudaLaunchKernel(const void *func, dim3 gridDim, dim3 blockDim, void **args, size_t sharedMem, cudaStream_t stream) {
        // real target func
        static cudaLaunchKernel_t real_cudaLaunchKernel = nullptr;
        auto tid = gettid();

        if (!real_cudaLaunchKernel) {
            auto lib_path = get_shared_library_path("libcudart.so");
            real_cudaLaunchKernel = load_func<cudaLaunchKernel_t>("cudaLaunchKernel", lib_path);
            if(!real_cudaLaunchKernel) return cudaErrorInvalidValue;
        }
        if(isInContext("cudaLaunchKernel")){
            std::ostringstream oss;
            auto mangle_func_dec = load_func_name(func);

            std::string func_name, params_str;
            get_params(mangle_func_dec, args, func_name, params_str);
            std::string extra_msg = getEnv("CUDA_API_HOOK_EXTRA");
            oss << " [" << extra_msg << "] " << func_name << "\n  <<<(" << gridDim.x << "," << gridDim.y << "," << gridDim.z << "),"
                << "(" << blockDim.x << "," << blockDim.y << "," << blockDim.z << "),(" << sharedMem << "),(" 
                << size_t(stream) << ")>>>\n  " << params_str << std::endl;
            std::cout << HOOK_LOG_PREFIX << oss.str() << std::endl;
        }
        cudaError_t err = real_cudaLaunchKernel(func, gridDim, blockDim, args, sharedMem, stream); 
        return err;
    }

}// end extern "C"

/*ref cmd

g++ -shared -fPIC -o cuda_api_hook.so cuda_api_hook.cpp -ldl -I/usr/local/cuda-12.8/targets/x86_64-linux/include/
export LD_PRELOAD=/your_so_dir/cuda_api_hook.so

g++ -o cuda_api_hook cuda_api_hook.cpp -ldl -I/usr/local/cuda-12.8/targets/x86_64-linux/include/

*/

//test
int main(){
    std::vector<std::string> cases = {
        "_ZN4vllm30reshape_and_cache_flash_kernelI13__nv_bfloat16S1_LNS_18Fp8KVCacheDataTypeE0EEEvPKT_S5_PT0_S7_PKlllllliiiPKfSB_",
        "_ZN2at6native29vectorized_elementwise_kernelILi4ENS0_11FillFunctorIiEESt5arrayIPcLm1EEEEviT0_T1_",
        "_Z33append_decode_cache_T_rope_kernelI13__nv_bfloat16Li8EEvPKT_PS1_S4_S4_PKiS6_S6_S6_S6_PKfS8_iiiiijib", 
        "_Z43multi_query_append_attention_warp1_4_kernelI13__nv_bfloat16Lb0ELj5ELb1ELj4ELj1ELj4ELj128ELj64ELj1ELj1ELj8ES0_Lb1EEvPT_S2_S2_PKS1_S4_PKiS6_S6_S6_S6_S6_iiiffffjS2_PfS7_PT11_i"
    };
    for(auto &name : cases){
        std::string demangle_func_dec = demangle(name), func_name, args_decl;
        extract_name_and_args(demangle_func_dec, func_name, args_decl);
        std::cout << demangle_func_dec << "\n"
            << func_name << "\n"
            << args_decl << "\n";
        for(auto &param_info : parseFuncParams(args_decl)){
            std::cout << "(" << param_info.first << "," << DataTypeToStr(param_info.second) << "), ";
        }
        std::cout << std::endl;
        std::cout << "\n";
    }
    return 0;
}